# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: site-packages\ntlm_auth\messages.py
import hmac, os, struct
from ntlm_auth.compute_response import ComputeResponse
from ntlm_auth.constants import NegotiateFlags, MessageTypes, NTLM_SIGNATURE, AvFlags
from ntlm_auth.rc4 import ARC4
from ntlm_auth.target_info import TargetInfo


class NegotiateMessage(object):
    EXPECTED_BODY_LENGTH = 40

    def __init__(self, negotiate_flags, domain_name, workstation):
        self.signature = NTLM_SIGNATURE
        self.message_type = struct.pack("<L", MessageTypes.NTLM_NEGOTIATE)
        if domain_name is None:
            self.domain_name = ""
        else:
            self.domain_name = domain_name
            negotiate_flags |= NegotiateFlags.NTLMSSP_NEGOTIATE_OEM_DOMAIN_SUPPLIED
        if workstation is None:
            self.workstation = ""
        else:
            self.workstation = workstation
            negotiate_flags |= NegotiateFlags.NTLMSSP_NEGOTIATE_OEM_WORKSTATION_SUPPLIED
        negotiate_flags -= NegotiateFlags.NTLMSSP_NEGOTIATE_UNICODE
        negotiate_flags |= NegotiateFlags.NTLMSSP_NEGOTIATE_OEM
        self.domain_name = self.domain_name.encode("ascii")
        self.workstation = self.workstation.encode("ascii")
        self.version = get_version(negotiate_flags)
        self.negotiate_flags = struct.pack("<I", negotiate_flags)

    def get_data(self):
        payload_offset = self.EXPECTED_BODY_LENGTH
        domain_name_len = struct.pack("<H", len(self.domain_name))
        domain_name_max_len = struct.pack("<H", len(self.domain_name))
        domain_name_buffer_offset = struct.pack("<I", payload_offset)
        payload_offset += len(self.domain_name)
        workstation_len = struct.pack("<H", len(self.workstation))
        workstation_max_len = struct.pack("<H", len(self.workstation))
        workstation_buffer_offset = struct.pack("<I", payload_offset)
        payload_offset += len(self.workstation)
        payload = self.domain_name
        payload += self.workstation
        msg1 = self.signature
        msg1 += self.message_type
        msg1 += self.negotiate_flags
        msg1 += domain_name_len
        msg1 += domain_name_max_len
        msg1 += domain_name_buffer_offset
        msg1 += workstation_len
        msg1 += workstation_max_len
        msg1 += workstation_buffer_offset
        msg1 += self.version
        if not self.EXPECTED_BODY_LENGTH == len(msg1):
            raise AssertionError(
                "BODY_LENGTH: %d != msg1: %d" % (self.EXPECTED_BODY_LENGTH, len(msg1))
            )
        msg1 += payload
        return msg1


class ChallengeMessage(object):
    """
        [MS-NLMP] v28.0 2016-07-14
    
        2.2.1.2 CHALLENGE_MESSAGE
        The CHALLENGE_MESSAGE defines an NTLM challenge message that is sent from the server to
        the client. The CHALLENGE_MESSAGE is used by the server to challenge the client to prove
        its identity, For connection-oriented requests, the CHALLENGE_MESSAGE generated by the
        server is in response to the NEGOTIATE_MESSAGE from the client.
    
        :param msg2: The CHALLENGE_MESSAGE received from the server after sending our NEGOTIATE_MESSAGE. This has
                        been decoded from a base64 string
    
        Attributes
            signature: An 8-byte character array that MUST contain the ASCII string 'NTLMSSP\x00'
            message_type: A 32-bit unsigned integer that indicates the message type. This field must be set to 0x00000002
            negotiate_flags: A NEGOTIATE strucutre that contains a set of bit flags. The server sets flags to indicate options it supports
            server_challenge: A 64-bit value that contains the NTLM challenge. The challenge is a 64-bit nonce. Used in the AuthenticateMessage message
            reserved: An 8-byte array whose elements MUST be zero when sent and MUST be ignored on receipt
            version: When NTLMSSP_NEGOTIATE_VERSION flag is set in negotiate_flags field which contains the windows version info. Used only for debugging purposes
            target_name: When NTLMSSP_REQUEST_TARGET is set is a byte array that contains the name of the server authentication realm. In a domain environment this is the domain name not server name
            target_info: When NTLMSSP_NEGOTIATE_TARGET_INFO is set is a byte array that contains a sequence of AV_PAIR structures (target_info.py)
    """

    def __init__(self, msg2):
        self.data = msg2
        self.signature = msg2[0:8]
        self.message_type = struct.unpack("<I", msg2[8:12])[0]
        self.negotiate_flags = struct.unpack("<I", msg2[20:24])[0]
        self.server_challenge = msg2[24:32]
        self.reserved = msg2[32:40]
        if self.negotiate_flags & NegotiateFlags.NTLMSSP_NEGOTIATE_VERSION:
            if (
                self.negotiate_flags
                & NegotiateFlags.NTLMSSP_NEGOTIATE_EXTENDED_SESSIONSECURITY
            ):
                if len(msg2) > 48:
                    self.version = struct.unpack("<q", msg2[48:56])[0]
                self.version = None
            if self.negotiate_flags & NegotiateFlags.NTLMSSP_REQUEST_TARGET:
                target_name_len = struct.unpack("<H", msg2[12:14])[0]
                target_name_max_len = struct.unpack("<H", msg2[14:16])[0]
                target_name_buffer_offset = struct.unpack("<I", msg2[16:20])[0]
                self.target_name = msg2[
                    target_name_buffer_offset : target_name_buffer_offset
                    + target_name_len
                ]
            else:
                self.target_name = None
            if self.negotiate_flags & NegotiateFlags.NTLMSSP_NEGOTIATE_TARGET_INFO:
                target_info_len = struct.unpack("<H", msg2[40:42])[0]
                target_info_max_len = struct.unpack("<H", msg2[42:44])[0]
                target_info_buffer_offset = struct.unpack("<I", msg2[44:48])[0]
                target_info_raw = msg2[
                    target_info_buffer_offset : target_info_buffer_offset
                    + target_info_len
                ]
                self.target_info = TargetInfo(target_info_raw)
            else:
                self.target_info = None
            if not self.signature == NTLM_SIGNATURE:
                raise AssertionError
            if not self.message_type == MessageTypes.NTLM_CHALLENGE:
                raise AssertionError

    def get_data(self):
        return self.data


class AuthenticateMessage(object):
    EXPECTED_BODY_LENGTH = 72
    EXPECTED_BODY_LENGTH_WITH_MIC = 88

    def __init__(
        self,
        user_name,
        password,
        domain_name,
        workstation,
        challenge_message,
        ntlm_compatibility,
        server_certificate_hash,
    ):
        self.signature = NTLM_SIGNATURE
        self.message_type = struct.pack("<L", MessageTypes.NTLM_AUTHENTICATE)
        self.negotiate_flags = challenge_message.negotiate_flags
        self.version = get_version(self.negotiate_flags)
        self.mic = None
        if domain_name is None:
            self.domain_name = ""
        else:
            self.domain_name = domain_name
        if workstation is None:
            self.workstation = ""
        else:
            self.workstation = workstation
        if self.negotiate_flags & NegotiateFlags.NTLMSSP_NEGOTIATE_UNICODE:
            self.negotiate_flags -= NegotiateFlags.NTLMSSP_NEGOTIATE_OEM
            encoding_value = "utf-16-le"
        else:
            encoding_value = "ascii"
        self.domain_name = self.domain_name.encode(encoding_value)
        self.user_name = user_name.encode(encoding_value)
        self.workstation = self.workstation.encode(encoding_value)
        compute_response = ComputeResponse(
            user_name, password, domain_name, challenge_message, ntlm_compatibility
        )
        self.lm_challenge_response = compute_response.get_lm_challenge_response()
        self.nt_challenge_response, key_exchange_key, target_info = compute_response.get_nt_challenge_response(
            self.lm_challenge_response, server_certificate_hash
        )
        self.target_info = target_info
        if self.negotiate_flags & NegotiateFlags.NTLMSSP_NEGOTIATE_KEY_EXCH:
            self.exported_session_key = get_random_export_session_key()
            rc4_handle = ARC4(key_exchange_key)
            self.encrypted_random_session_key = rc4_handle.update(
                self.exported_session_key
            )
        else:
            self.exported_session_key = key_exchange_key
            self.encrypted_random_session_key = b""
        self.negotiate_flags = struct.pack("<I", self.negotiate_flags)

    def get_data(self):
        if self.mic is None:
            mic = b""
            expected_body_length = self.EXPECTED_BODY_LENGTH
        else:
            mic = self.mic
            expected_body_length = self.EXPECTED_BODY_LENGTH_WITH_MIC
        payload_offset = expected_body_length
        domain_name_len = struct.pack("<H", len(self.domain_name))
        domain_name_max_len = struct.pack("<H", len(self.domain_name))
        domain_name_buffer_offset = struct.pack("<I", payload_offset)
        payload_offset += len(self.domain_name)
        user_name_len = struct.pack("<H", len(self.user_name))
        user_name_max_len = struct.pack("<H", len(self.user_name))
        user_name_buffer_offset = struct.pack("<I", payload_offset)
        payload_offset += len(self.user_name)
        workstation_len = struct.pack("<H", len(self.workstation))
        workstation_max_len = struct.pack("<H", len(self.workstation))
        workstation_buffer_offset = struct.pack("<I", payload_offset)
        payload_offset += len(self.workstation)
        lm_challenge_response_len = struct.pack("<H", len(self.lm_challenge_response))
        lm_challenge_response_max_len = struct.pack(
            "<H", len(self.lm_challenge_response)
        )
        lm_challenge_response_buffer_offset = struct.pack("<I", payload_offset)
        payload_offset += len(self.lm_challenge_response)
        nt_challenge_response_len = struct.pack("<H", len(self.nt_challenge_response))
        nt_challenge_response_max_len = struct.pack(
            "<H", len(self.nt_challenge_response)
        )
        nt_challenge_response_buffer_offset = struct.pack("<I", payload_offset)
        payload_offset += len(self.nt_challenge_response)
        encrypted_random_session_key_len = struct.pack(
            "<H", len(self.encrypted_random_session_key)
        )
        encrypted_random_session_key_max_len = struct.pack(
            "<H", len(self.encrypted_random_session_key)
        )
        encrypted_random_session_key_buffer_offset = struct.pack("<I", payload_offset)
        payload_offset += len(self.encrypted_random_session_key)
        payload = self.domain_name
        payload += self.user_name
        payload += self.workstation
        payload += self.lm_challenge_response
        payload += self.nt_challenge_response
        payload += self.encrypted_random_session_key
        msg3 = self.signature
        msg3 += self.message_type
        msg3 += (
            lm_challenge_response_len
            + lm_challenge_response_max_len
            + lm_challenge_response_buffer_offset
        )
        msg3 += (
            nt_challenge_response_len
            + nt_challenge_response_max_len
            + nt_challenge_response_buffer_offset
        )
        msg3 += domain_name_len + domain_name_max_len + domain_name_buffer_offset
        msg3 += user_name_len + user_name_max_len + user_name_buffer_offset
        msg3 += workstation_len + workstation_max_len + workstation_buffer_offset
        msg3 += (
            encrypted_random_session_key_len
            + encrypted_random_session_key_max_len
            + encrypted_random_session_key_buffer_offset
        )
        msg3 += self.negotiate_flags
        msg3 += self.version
        msg3 += mic
        msg3 += payload
        return msg3

    def add_mic(self, negotiate_message, challenge_message):
        if self.target_info is not None:
            av_flags = self.target_info[TargetInfo.MSV_AV_FLAGS]
            if av_flags is not None:
                if av_flags[1] == struct.pack("<L", AvFlags.MIC_PROVIDED):
                    self.mic = struct.pack("<IIII", 0, 0, 0, 0)
                    negotiate_data = negotiate_message.get_data()
                    challenge_data = challenge_message.get_data()
                    authenticate_data = self.get_data()
                    mic = hmac.new(
                        self.exported_session_key,
                        negotiate_data + challenge_data + authenticate_data,
                    ).digest()
                    self.mic = mic


def get_version(negotiate_flags):
    if negotiate_flags & NegotiateFlags.NTLMSSP_NEGOTIATE_VERSION:
        product_major_version = struct.pack("<B", 6)
        product_minor_version = struct.pack("<B", 1)
        product_build = struct.pack("<H", 7601)
        version_reserved = b"\x00\x00\x00"
        ntlm_revision_current = struct.pack("<B", 15)
        version = (
            product_major_version
            + product_minor_version
            + product_build
            + version_reserved
            + ntlm_revision_current
        )
    else:
        version = b"\x00\x00\x00\x00\x00\x00\x00\x00"
    return version


def get_random_export_session_key():
    return os.urandom(16)
